package com.my.study.limiter;

import com.google.common.math.LongMath;
import com.my.study.util.RedisUtil;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RLock;

import java.util.concurrent.TimeUnit;

/**
 * @author: Carlos
 * @create: 2020/12/22 17:00
 */
@Slf4j
@Setter
@Getter
public class RedisLimiter {

    /**
     * redis key
     */
    private String key;

    /**
     * 每秒存入的令牌数
     */
    private Double permitsPerSecond;

    /**
     * 最大maxBurstSeconds秒生成的令牌
     */
    private Integer maxBurstSeconds;

    private RLock lock;

    /**
     * redis工具类
     */
    private RedisUtil redisUtil;

    public RedisLimiter(String key, Double permitsPerSecond, Integer maxBurstSeconds, RLock lock, RedisUtil redisUtil) {
        this.key = key;
        this.permitsPerSecond = permitsPerSecond;
        this.maxBurstSeconds = maxBurstSeconds;
        this.lock = lock;
        this.redisUtil = redisUtil;
    }

    /**
     * 生成并存储默认令牌桶
     *
     * @return RedisPermits
     * @author Carlos
     * @date 2020/12/28 15:30
     */
    private RedisPermits putDefaultPermits() {
        lock.lock();
        try {
            Object obj = redisUtil.get(key);
            if (null == obj) {
                RedisPermits permits = new RedisPermits(permitsPerSecond, maxBurstSeconds);
                redisUtil.set(key, permits, permits.expires());
                return permits;
            } else {
                return ((RedisPermits) obj);
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 获取令牌桶
     *
     * @return RedisPermits
     * @author Carlos
     * @date 2020/12/28 15:40
     */
    public RedisPermits getPermits() {
        Object obj = redisUtil.get(key);
        if (null == obj) {
            return this.putDefaultPermits();
        }
        return ((RedisPermits) obj);
    }

    /**
     * 设置令牌桶
     *
     * @param redisPermits 令牌桶
     * @author Carlos
     * @date 2020/12/28 15:50
     */
    public void setPermits(RedisPermits redisPermits) {
        redisUtil.set(key, redisPermits, redisPermits.expires());
    }

    /**
     * 等待直到获取指定数量的令牌
     *
     * @param tokens 指定令牌数量
     * @return
     */
    public long acquire(long tokens) throws InterruptedException {
        long milliToWait = this.reserve(tokens);
        log.info("acquire for {}ms {}", milliToWait, Thread.currentThread().getName());
        Thread.sleep(milliToWait);
        return milliToWait;
    }

    /**
     * 获取1个令牌数
     *
     * @return
     * @throws InterruptedException
     */
    public long acquire() throws InterruptedException {
        return acquire(1L);
    }

    /**
     * 尝试获取令牌
     *
     * @param tokens  要获取的令牌数
     * @param timeout  获取令牌等待的时间，最小为0
     * @param unit
     * @return
     * @throws InterruptedException
     */
    private boolean tryAcquire(long tokens, long timeout, TimeUnit unit) throws InterruptedException {
        long timeoutMicros = Math.max(unit.toMillis(timeout), 0);
        this.checkTokens(tokens);
        Long milliToWait;
        lock.lock();
        try {
            if (!this.canAcquire(tokens, timeoutMicros)) {
                return false;
            } else {
                milliToWait = this.reserveAndGetWaitLength(tokens);
            }
        } finally {
            lock.unlock();
        }
        Thread.sleep(milliToWait);
        return true;
    }

    /**
     * 获取一个令牌
     * @param timeout
     * @param unit
     * @return
     * @throws InterruptedException
     */
    private Boolean tryAcquire(Long timeout , TimeUnit unit) throws InterruptedException{
        return tryAcquire(1L,timeout, unit);
    }

    /**
     * 在等待的时间内是否可以获取到令牌
     *
     * @param tokens
     * @param timeoutMicros
     * @return
     */
    private boolean canAcquire(long tokens, long timeoutMicros) {
        return queryEarliestAvailable(tokens) - timeoutMicros <= 0;
    }

    /**
     * 获取指定令牌所需等待的时间
     *
     * @param tokens
     * @return 等待的时间
     *
     */
    private long reserve(long tokens) {
        this.checkTokens(tokens);
        lock.lock();
        try {
            return this.reserveAndGetWaitLength(tokens);
        } finally {
            lock.unlock();
        }
    }

    /**
     * 校验token值 token值需要为正整数(>0)
     *
     * @param tokens
     */
    private void checkTokens(long tokens) {
        if (tokens < 0) {
            throw new IllegalArgumentException("Requested tokens " + tokens + " must be positive");
        }
    }

    /**
     * 返回获取{tokens}个令牌最早可用的时间
     *
     * @param tokens
     * @return
     */
    private long queryEarliestAvailable(long tokens) {
        long n = System.currentTimeMillis();
        RedisPermits permits = this.getPermits();
        permits.reSync(n);
        // 可以消耗的令牌数
        long storedPermitsToSpend = Math.min(tokens, permits.getStoredPermits());
        // 需要等待的令牌数
        long freshPermits = tokens - storedPermitsToSpend;
        // 需要等待的时间
        long waitMillis = freshPermits * permits.getNextFreeTicketMillis();
        return LongMath.saturatedAdd(permits.getNextFreeTicketMillis() - n, waitMillis);
    }

    /**
     * 预定tokens个令牌并返回所需要等待的时间
     *
     * @param tokens
     * @return
     */
    private long reserveAndGetWaitLength(long tokens) {
        long n = System.currentTimeMillis();
        RedisPermits permits = this.getPermits();
        permits.reSync(n);
        // 可以消耗的令牌数
        long storedPermitsToSpend  = Math.min(tokens, permits.getStoredPermits());
        // 需要等待的令牌数
        long freshPermits  = tokens - storedPermitsToSpend;
        // 需要等待的时间
        long waitMillis = freshPermits * permits.getIntervalMillis();
        permits.setNextFreeTicketMillis(LongMath.saturatedAdd(permits.getNextFreeTicketMillis(), waitMillis));
        permits.setStoredPermits(permits.getStoredPermits() - storedPermitsToSpend);
        this.setPermits(permits);
        return permits.getNextFreeTicketMillis() - n;
    }

}
